package com.zyx.flinkdemo.sql.func;

import com.zyx.flinkdemo.pojo.Top2Accum;
import com.zyx.flinkdemo.pojo.WaterSensor;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.table.api.Table;
import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
import org.apache.flink.table.functions.TableAggregateFunction;
import org.apache.flink.util.Collector;

import static org.apache.flink.table.api.Expressions.$;
import static org.apache.flink.table.api.Expressions.call;

/**
 * @author Yaxi.Zhang
 * @since 2021/5/31 13:45
 * desc: UDTAF(表值聚合函数)案例
 */
public class UdtafDemo {
    public static void main(String[] args) throws Exception {
        // 1.获取执行环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        StreamTableEnvironment tableEnv = StreamTableEnvironment.create(env);

        // 2.读取端口数据并转换为JavaBean
        SingleOutputStreamOperator<WaterSensor> waterSensorDs = env
                .socketTextStream("localhost", 7777)
                .map(line -> {
                    String[] split = line.split(",");
                    return new WaterSensor(split[0],
                            Long.parseLong(split[1]),
                            Integer.parseInt(split[2]));
                });

        // 3.将流转换为动态表
        Table table = tableEnv.fromDataStream(waterSensorDs);

        // 4.先注册再使用
        tableEnv.createTemporarySystemFunction("top2vc", Top2Vc.class);
        tableEnv.createTemporaryView("water_sensor", table);

        //TableAPI
        table
                .groupBy($("id"))
                // Deprecated
                .flatAggregate(call("top2vc", $("vc")).as("top", "rank"))
                // Deprecated
                // .flatAggregate("top2vc(vc) as (top, rank)")
                .select($("id"), $("top"), $("rank"))
                .execute()
                .print();

        //6.执行任务
        env.execute();

    }

    public static class Top2Vc extends TableAggregateFunction<Tuple2<Integer, String>, Top2Accum> {

        @Override
        public Top2Accum createAccumulator() {
            return new Top2Accum(Integer.MIN_VALUE, Integer.MIN_VALUE);
        }

        public void accumulate(Top2Accum acc, Integer v) {
            if (v > acc.first) {
                acc.second = acc.first;
                acc.first = v;
            } else if (v > acc.second) {
                acc.second = v;
            }
        }

        public void emitValue(Top2Accum acc, Collector<Tuple2<Integer, String>> out) {
            // emit the value and rank
            if (acc.first != Integer.MIN_VALUE) {
                out.collect(Tuple2.of(acc.first, "top1"));
            }
            if (acc.second != Integer.MIN_VALUE) {
                out.collect(Tuple2.of(acc.second, "top2"));
            }
        }
    }
}
